import open3d as o3d
import numpy as np
import matplotlib.pyplot as plt

# 参数配置
pcd_path = "../maps/map.pcd"
output_path = "../maps/stair1_clip.pcd"
min_bound = np.array([-10, -4, -5], dtype=np.float64)
max_bound = np.array([7, 8, 20], dtype=np.float64)
down_sample_resolution = 0.01

# 1. 读取点云
pcd = o3d.io.read_point_cloud(pcd_path)
print(f"原始点数: {np.asarray(pcd.points).shape[0]}")

# 2. 边界打印
bbox = pcd.get_axis_aligned_bounding_box()
print(f"边界：min={bbox.min_bound}, max={bbox.max_bound}")

# 3. 降采样
pcd_down = pcd.voxel_down_sample(down_sample_resolution)
points = np.asarray(pcd_down.points)
print(f"降采样后点数: {points.shape[0]}")

# 4. 有效点提取
valid_mask = (
    (points[:, 0] >= min_bound[0])
    & (points[:, 0] <= max_bound[0])
    & (points[:, 1] >= min_bound[1])
    & (points[:, 1] <= max_bound[1])
    & (points[:, 2] >= min_bound[2])
    & (points[:, 2] <= max_bound[2])
)
print(f"有效点数: {np.count_nonzero(valid_mask)}")

# 5. 着色
colors = np.zeros((points.shape[0], 3), dtype=np.float64)  # 确保为float64
z_vals = points[valid_mask][:, 2]
if len(z_vals) > 0:
    z_min, z_max = z_vals.min(), z_vals.max()
    # 每2米一个循环
    cycle_length = 2.0
    z_cycle = (z_vals - z_min) % cycle_length
    z_norm = z_cycle / cycle_length
    cmap = plt.get_cmap("hsv")
    color_map = cmap(z_norm)[:, :3].astype(np.float64)  # 显式转换
    colors[valid_mask] = color_map
else:
    print("警告：未找到有效点！")

pcd_down.colors = o3d.utility.Vector3dVector(colors)

# 6. 保存有效点（仅有效点）
valid_pcd = o3d.geometry.PointCloud()
valid_pcd.points = o3d.utility.Vector3dVector(points[valid_mask])
valid_pcd.colors = o3d.utility.Vector3dVector(colors[valid_mask])
o3d.io.write_point_cloud(output_path, valid_pcd)
print(f"保存有效点云至: {output_path}")

# 提取无效点
invalid_pcd = o3d.geometry.PointCloud()
invalid_pcd.points = o3d.utility.Vector3dVector(points[~valid_mask])
invalid_pcd.colors = o3d.utility.Vector3dVector(colors[~valid_mask])

# 可视化
o3d.visualization.draw_geometries(
    [valid_pcd, invalid_pcd],
    window_name="有效区域点云（彩色）",
    width=1920,
    height=1080,
    left=50,
    top=50,
    point_show_normal=False,
    mesh_show_wireframe=False,
    mesh_show_back_face=False
)
